import { RenderTexture, Texture2D } from "cc";
import { BuiltInSensorType, IBuiltInSensor } from "./IBuiltInSensor";
import { ISensor } from "./ISensor";
import { ObservationSpec } from "./ObservationSpec";
import { CompressionSpec, SensorCompressionType } from "./CompressionSpec";
import { ObservationWriter, ObservationWriterExtension } from "./ObservationWriter";
import { Texture2DWrapper } from "../Utils/Texture2DWrapper";
import { Utilities } from "../Utilities";


export class RenderTextureSensor implements ISensor, IBuiltInSensor {
    private _renderTexture: RenderTexture;
    private _grayscale: boolean;
    private _name: string;
    private _observationSpec: ObservationSpec;
    private _compressionType: SensorCompressionType;
    private _texture: Texture2D = null!;

    public get compressionType(): SensorCompressionType {
        return this._compressionType;
    }

    public set compressionType(value: SensorCompressionType) {
        this._compressionType = value;
    }

    /**
     * 
     * @param renderTexture 
     * @param grayscale whether to convert it to grayscale or not 
     * @param name  name of the sensor
     * @param compressionType  compression method for the render texture
     */
    constructor(renderTexture: RenderTexture, grayscale: boolean, name: string, compressionType: SensorCompressionType) {
        this._renderTexture = renderTexture;
        const width = renderTexture ? renderTexture.width : 0;
        const height = renderTexture ? renderTexture.height : 0;
        this._grayscale = grayscale;
        this._name = name;
        this._observationSpec = ObservationSpec.visual(height, width, grayscale ? 1 : 3);
        this._compressionType = compressionType;


        this._texture = new Texture2D();
        this._texture.reset({
            width: width,
            height: height,
            format: Texture2D.PixelFormat.RGBA8888,
            mipmapLevel: 1,
            baseLevel: 0,
            maxLevel: 100
        })
    }

    public getName(): string {
        return this._name;
    }

    public getObservationSpec(): ObservationSpec {
        return this._observationSpec;
    }

    public getCompressedObservation(): Uint8Array {
        RenderTextureSensor.observationToTexture(this._renderTexture, this._texture);
        const compressed = Texture2DWrapper.encodeToPNG(this._texture) ?? new Uint8Array();
        return compressed;
    }

    public write(writer: ObservationWriter): number {
        RenderTextureSensor.observationToTexture(this._renderTexture, this._texture);
        const numWritten = ObservationWriterExtension.writeTexture(writer, this._texture, this._grayscale);
        return numWritten;
    }

    public update(): void { }

    public reset(): void { }

    public getCompressionSpec(): CompressionSpec {
        return new CompressionSpec(this._compressionType);
    }

    public getBuiltInSensorType(): BuiltInSensorType {
        return BuiltInSensorType.RenderTextureSensor;
    }

    public dispose(): void {
        if (this._texture) {
            Utilities.destroyTexture(this._texture);
            this._texture = null!;
        }
    }

    private static observationToTexture(obsTexture: RenderTexture, texture2D: Texture2D): void {
        const pixels = obsTexture.readPixels();
        if (pixels) {
            texture2D.uploadData(pixels);
        }
    }
}